#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <napi.h>
#include <iostream>
#include "stb_image.h"
#include "stb_image_write.h"

using namespace Napi;


void rgb888_to_rgb565(uint8_t r, uint8_t g, uint8_t b, uint8_t *result) {
  uint16_t red = r >> 3;  // 5位红
  uint16_t green = g >> 2; // 6位绿
  uint16_t blue = b >> 3;  // 5位蓝
  uint16_t temp = (red << 11) | (green << 5) | blue;

  result[0] = temp & 0xFF;
  result[1] = temp >> 8;
}

Napi::Value image_to_matrix_path(const Napi::CallbackInfo &info)
{
  Napi::Env env = info.Env();
  std::string path = info[0].As<Napi::String>().Utf8Value();
  std::string type = info[1].As<Napi::String>().Utf8Value();

  int width, height, channels;
  unsigned char *image_data = stbi_load(path.c_str(), &width, &height, &channels, 0);

  if (type == "RGB")
  {
    Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 3);

    if (channels == 4)
    {
      std::vector<uint8_t> buffer;
      for (long i = 0; i < width * height; i++)
      {
        if (i == 0 || i % 3 != 0)
        {
          buffer.push_back(image_data[i]);
        }
      }

      memcpy(result.Data(), buffer.data(), width * height * 3);
      stbi_image_free(image_data);
      return result;
    }
    else if (channels != 3)
    {
      stbi_image_free(image_data);
      Napi::TypeError::New(env, "Invalid channels").ThrowAsJavaScriptException();
      return env.Null();
    }

    memcpy(result.Data(), image_data, width * height * 3);
    stbi_image_free(image_data);
    return result;
  }
  else if (type == "RGBA")
  {
    Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 4);
    if (channels == 4)
    {
      memcpy(result.Data(), image_data, width * height * 4);
      stbi_image_free(image_data);
      return result;
    }
    stbi_image_free(image_data);
    Napi::TypeError::New(env, "Invalid channels").ThrowAsJavaScriptException();
    return env.Null();
  }
  else if (type == "RGB565")
  {
    if (channels == 4)
    {
      Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 2);
      std::vector<uint8_t> buffer;
      for (long i = 0; i < width * height; i++)
      {
        if (i == 0 || i % 3 != 0)
        {
          buffer.push_back(image_data[i]);
        }
      }

      std::vector<uint8_t> data;
      for (long i = 0; i < width * height; i += 3)
      {
        uint8_t temp[2];
        rgb888_to_rgb565(buffer[i], buffer[i + 1], buffer[i + 2], temp);
        data.push_back(temp[0]);
        data.push_back(temp[1]);
      }

      memcpy(result.Data(), buffer.data(), width * height * 2);
      stbi_image_free(image_data);
      return result;
    }
    else if (channels == 3)
    {
      Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 2);
      std::vector<uint8_t> buffer;
      for (long i = 0; i < width * height * 3; i += 3 )
      {
        uint8_t temp[2] = {0, 0};
        rgb888_to_rgb565(image_data[i], image_data[i + 1], image_data[i + 2], temp);
        buffer.push_back(temp[0]);
        buffer.push_back(temp[1]);
      }

      memcpy(result.Data(), buffer.data(), width * height * 2);
      stbi_image_free(image_data);
      return result;
    }
    stbi_image_free(image_data);
    Napi::TypeError::New(env, "Invalid channels").ThrowAsJavaScriptException();
    return env.Null();
  }
  else
  {
    stbi_image_free(image_data);
    Napi::TypeError::New(env, "Invalid type").ThrowAsJavaScriptException();
    return env.Null();
  }
}

Napi::Value image_to_matrix_buffer(const Napi::CallbackInfo &info)
{
  Napi::Env env = info.Env();
  Napi::Buffer<uint8_t> data = info[0].As<Napi::Buffer<uint8_t>>();
  std::string type = info[1].As<Napi::String>().Utf8Value();

  int width, height, channels;
  unsigned char *image_data = stbi_load_from_memory((stbi_uc *)data.Data(), data.Length(), &width, &height, &channels, 0);


  if (type == "RGB") {
    Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 3);

    if (channels == 4) {
      std::vector<uint8_t> buffer;
      for (long i = 0; i < width * height; i++) {
        if (i == 0 || i % 3 != 0) {
          buffer.push_back(image_data[i]);
        }
      }

      memcpy(result.Data(), buffer.data(), width * height * 3);
      stbi_image_free(image_data);
      return result;
    } else if (channels != 3) {
      stbi_image_free(image_data);
      Napi::TypeError::New(env, "Invalid channels").ThrowAsJavaScriptException();
      return env.Null();
    }

    memcpy(result.Data(), image_data, width * height * 3);
    stbi_image_free(image_data);
    return result;
  } else if (type == "RGBA") {
    Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 4);
    if (channels == 4) {
      memcpy(result.Data(), image_data, width * height * 4);
      stbi_image_free(image_data);
      return result;
    }
    stbi_image_free(image_data);
    Napi::TypeError::New(env, "Invalid channels").ThrowAsJavaScriptException();
    return env.Null();
  } else if (type == "RGB565") {
    Napi::Buffer<uint8_t> result = Napi::Buffer<uint8_t>::New(env, width * height * 2);
    if (channels == 4) {
      std::vector<uint8_t> buffer;
      for (long i = 0; i < width * height; i++)
      {
        if (i == 0 || i % 3 != 0)
        {
          buffer.push_back(image_data[i]);
        }
      }

      std::vector<uint8_t> data;
      for (long i = 0; i < width * height; i++)
      {
        uint8_t temp[2];
        rgb888_to_rgb565(buffer[i], buffer[i + 1], buffer[i + 2], temp);
        data.push_back(temp[0]);
        data.push_back(temp[1]);
      }

      memcpy(result.Data(), buffer.data(), width * height * 2);
      stbi_image_free(image_data);
      return result;
    }
    else if(channels == 3)
    {
      std::vector<uint8_t> data;
      for (long i = 0; i < width * height; i++)
      {
        uint8_t temp[2];
        rgb888_to_rgb565(image_data[i], image_data[i + 1], image_data[i + 2], temp);
        data.push_back(temp[0]);
        data.push_back(temp[1]);
      }

      memcpy(result.Data(), image_data, width * height * 2);
      stbi_image_free(image_data);
      return result;
    }
    stbi_image_free(image_data);
    Napi::TypeError::New(env, "Invalid channels").ThrowAsJavaScriptException();
    return env.Null();
  } else {
    stbi_image_free(image_data);
    Napi::TypeError::New(env, "Invalid type").ThrowAsJavaScriptException();
    return env.Null();
  }
}

Napi::Object Init(Napi::Env env, Napi::Object exports) {
  exports.Set(Napi::String::New(env, "imageToMatrixBuffer"),
              Napi::Function::New(env, image_to_matrix_buffer));
  exports.Set(Napi::String::New(env, "imageToMatrixPath"),
              Napi::Function::New(env, image_to_matrix_path));
  return exports;
}

NODE_API_MODULE(addon, Init)
